#include <iostream>
#include <coroutine>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <vector>
#include <queue>
#include <chrono>
#include <random>
#include "thread_pool.hpp"

using namespace std::chrono_literals;

// 协程任务类
class Task {
public:
    struct promise_type {
        Task get_return_object() { return {}; }
        std::suspend_never initial_suspend() { return {}; }
        std::suspend_never final_suspend() noexcept { return {}; }
        void return_void() {}
        void unhandled_exception() {}
    };
};

// 协程调度器
class Scheduler {
    std::mutex mtx;
    std::condition_variable cv;
    std::queue<std::coroutine_handle<>> ready_queue;
    bool stopped = false;

public:
    // 提交协程到调度器
    void submit(std::coroutine_handle<> h) {
        std::lock_guard<std::mutex> lock(mtx);
        ready_queue.push(h);
        cv.notify_one();
    }

    // 运行调度器
    void run() {
        while (true) {
            std::coroutine_handle<> h;
            {
                std::unique_lock<std::mutex> lock(mtx);
                cv.wait(lock, [this] {
                    return stopped || !ready_queue.empty();
                });

                if (stopped && ready_queue.empty()) break;

                h = ready_queue.front();
                ready_queue.pop();
            }

            if (h) {
                h.resume();
            }
        }
    }

    // 停止调度器
    void stop() {
        std::lock_guard<std::mutex> lock(mtx);
        stopped = true;
        cv.notify_all();
    }
};

Scheduler scheduler;

// 模拟阻塞操作
struct BlockingOp {
    static ThreadPool pool; // 静态线程池实例

    bool await_ready() { return false; }

    void await_suspend(std::coroutine_handle<> h) {
        pool.enqueue([h, this] {
            try {
                std::random_device rd;
                std::mt19937 gen(rd());
                std::uniform_int_distribution<> dis(100, 1000);
                std::this_thread::sleep_for(std::chrono::milliseconds(dis(gen)));
                std::cout << "BlockingOp done on thread " << std::this_thread::get_id() << std::endl;
                scheduler.submit(h); // 直接提交回调度器
            } catch(...) {
                std::cout << "BlockingOp exception on thread " << std::this_thread::get_id() << std::endl;
                scheduler.submit(h); // 确保协程不被卡死
            }
        });
    }

    void await_resume() {} // 不再需要条件变量等待
};

ThreadPool BlockingOp::pool(std::thread::hardware_concurrency());

// 协程任务函数
Task coroutine_task(int id) {
    co_await BlockingOp{}; // 立即挂起，确保首次恢复由调度器处理
    for (int i = 0; i < 3; ++i) {
        std::cout << "Task " << id << " step " << i << " on thread "
                  << std::this_thread::get_id() << std::endl;
        co_await BlockingOp{};
    }
}


int main() {
    // 初始任务

    for (int i = 0; i < 10; ++i) {
        coroutine_task(i);
        std::cout << "[初始任务] Task " << i << " 已加入" << " on thread "
                  << std::this_thread::get_id() << std::endl;
    }

    // 非阻塞运行调度器
    std::jthread scheduler_thread([&] { scheduler.run(); });

    // 模拟动态添加（每隔800ms加一个新任务）
    std::jthread dynamic_adder([] {
        for (int i = 10; i < 15; ++i) {
            std::this_thread::sleep_for(800ms);
            coroutine_task(i); // 动态添加任务
            std::cout << "[动态添加] Task " << i << " 已加入" << " on thread "
                  << std::this_thread::get_id() << std::endl;
        }
    });

    // 主线程等待
    std::this_thread::sleep_for(5s);
    scheduler.stop();
    return 0;
}
